"""
LangGraph Agent With Multi MCP Servers
"""

import asyncio
import json
import logging
import os
import sys
from typing import Any, Dict, List

from dotenv import load_dotenv
from langgraph.prebuilt import create_react_agent
from langchain_deepseek import ChatDeepSeek
from langchain_mcp_adapters.client import MultiServerMCPClient
from langgraph.checkpoint.memory import InMemorySaver
from mcp.server.fastmcp import FastMCP
from pathlib import Path

from insight_agent.utils.http_post import get_insight_request, inform_insight_thinking, update_insight_result
from langchain_qwq import ChatQwQ
# from tools.query_processing_server import query_processing
# from tools.text2sql_server import text_to_sql
# from tools.data_server import sql_inquire
# from tools.chart_server import visualize_sankey_chart, visualize_line_chart
from insight_agent.mcp_server_time_series.src.non_agent_mode import query_processing, text_to_sql, sql_inquire, visualize_line_chart, summary_factory, predict_future, sankey_to_sql, visualize_sankey_chart

# 添加utils路径到sys.path
sys.path.append(str(Path(__file__).resolve().parent.parent.parent / "utils"))

# 工作模式 react模式"react" 函数调用模式"non-agent"
mode = "non-agent"

# 构建文件读写路径
file_folder_path = Path().cwd().parent / "files"

# 初始化 MCP 服务器P
mcp = FastMCP("TimeSeriesServer", port=2129)

# 设置记忆存储
checkpointer = InMemorySaver()
current_dir = Path(__file__).resolve().parent  # 获取当前脚本的绝对路径

# 读取提示词
prompts_file_path = current_dir / "agent_time_series_prompts.txt"  # 构建绝对路径
with open(prompts_file_path, "r", encoding="utf-8") as f:
    prompt = f.read()

# 设置对话配置
config = {
    "configurable": {
        "thread_id": "1"
    }
}


class Configuration:
    """读取 .env 与 servers_config.json 等用以配置环境"""

    def __init__(self) -> None:
        load_dotenv()
        self.api_key: str = os.getenv("DASHSCOPE_API_KEY") or ""
        self.api_url: str = os.getenv("DASHSCOPE_API_BASE") or ""
        if not self.api_key:
            raise ValueError("❌ 未找到 LLM_API_KEY，请在 .env 中配置")

    servers_config_path = current_dir / "agent_time_series_servers_config.json"  # 构建绝对路径
    @staticmethod
    def load_servers(file_path: str = servers_config_path) -> Dict[str, Any]:
        with open(file_path, "r", encoding="utf-8") as f:
            return json.load(f).get("mcpServers", {})


async def run_chat_loop() -> None:
    """启动 MCP-Agent 聊天循环"""
    cfg = Configuration()
    # if cfg.base_url:
    #     os.environ["DEEPSEEK_API_BASE"] = cfg.base_url
    servers_cfg = Configuration.load_servers()

    # 连接多台 MCP 服务器
    mcp_client = MultiServerMCPClient(servers_cfg)
    tools = await mcp_client.get_tools()  # 工具对象列表
    logging.info(f"✅ 已加载 {len(tools)} 个 MCP 工具： {[t.name for t in tools]}")

    # 初始化大模型
    model = ChatQwQ(model="deepseek-v3")

    # 构造 LangGraph Agent
    agent = create_react_agent(model=model,
                               tools=tools,
                               prompt=prompt,
                               checkpointer=checkpointer)

    # CLI 多轮对话
    print("\n🤖 MCP Agent 已启动，输入 'quit' 退出")
    while True:
        user_input = input("\n你: ").strip()
        if user_input.lower() == "quit":
            break
        try:
            # # 非流式调用输出
            # result = await agent.ainvoke(
            #     {"messages": [{"role": "user", "content": user_input}]},
            #     config
            # )
            # print(f"\nAI: {result['messages'][-1].content}")

            async for step in agent.astream(
                    {"messages": [{"role": "user", "content": user_input}]}, config, stream_mode="values"
            ):
                step["messages"][-1].pretty_print()

            # # 流式调用输出
            # result = agent.stream(
            #     {"messages": [{"role": "user", "content": user_input}]},
            #     config,
            #     stream_mode="messages"
            # )
            # print(f"\nAI: ")
            # for message in result:
            #     print(message[0].content, end="")
        except Exception as exc:
            print(f"\n⚠️  出错: {exc}")

    # 清理
    # await mcp_client.cleanup()
    print("🧹 资源已清理，Bye!")

async def get_summary(query, response, data, agent):
    summary_prompt = f"""
    请扮演一个数据分析和总结撰写的角色，基于以下文本，生成一份关于用户问题和智能体分析结果和返回数据的分析总结。

    【背景信息】：
    - 用户的原始问题是 {query}
    
    【待总结的文本】：
    - 时序智能体的回答是 {response}
    - 时序智能体返回的图表数据是 {data}
    
    【具体要求】：
    1.  **重点突出**：请根据用户问题和智能体返回的文本和数据，进行有洞见的分析，抓住核心论点、结论和关键证据。
    2.  **结构清晰**：采用[如：总分总、列表、段落]的形式组织内容。
    3.  **长度控制**：将总结控制在[如：300字 / 4个要点]以内。
    4.  **语言风格**：使用[如：简洁、专业、口语化]的语言。
    5.  **输出格式**：请直接输出总结内容，以MD格式输出，不要添加额外解释。
    """

    ai_output = await agent.ainvoke(
        {"messages": [{"role": "user", "content": summary_prompt}]},
        config
    )

    return ai_output['messages'][-1].content

# 打包函数
def combine_info_to_post(summary: str, data: dict):
    updated_result =  {
        "summary": summary,
        "data": data
    }
    return updated_result

@mcp.tool()
async def agent_time_series(user_input):
    '''
    仅当用户输入与“时序、时间变化、流转关系、桑基图、分析XX与XX的流转关系”等相关的词语时使用本工具。
    本工具会自动解析用户请求，并返回分析报告。
    :param user_input: 用户请求
    :return:

    【适用场景】
    如果用户问题中如下关键字及其相似含义的短语，请强制调用本函数
    1. 近一周
    2. 近两个月
    3. 7天内
    4. 假设现在是YYYYMMDD，帮我预测7天内
    5. 分析流转关系
    6. 流转关系
    7. 其他与时间点和时间段相关的短语

    以下为适用的问题样例
    - 帮我分析20250717到20250730间北京市每日总家宽条数的变化
    - 请分析20250301之后7天的北京市人均本地通话费变化情况
    - 请分析在20250512-20250516期间的家庭圈个数变化情况
    - 请分析2023年1月至2023年12月的用户活跃度变化趋势
    - 帮我预测未来一周的流量峰值
    - 统计2024年5月每天的投诉量，并分析变化原因
    - 请分析20250701日 用户性别和全球通用户的流转关系
    - 用展示20250101日的省份和中高端用户的流转关系
    - 中高端用户与资费套餐的流转关系
    - 使用产品名称与所属集团的流转关系

    【返回内容】
    - 分析结果的 summary（简要结论）
    - 结构化的时序数据（如每日/每月的统计值、趋势线等）

    【注意事项】
    - 若用户未明确给出时间范围，可尝试智能补全或提示用户补充。
    - 仅在问题明确涉及“时间序列”或“趋势分析”时调用本工具，避免用于静态数据分析。

    '''
    import json
    if mode == "react":
        cfg = Configuration()
        # if cfg.base_url:
        #     os.environ["DEEPSEEK_API_BASE"] = cfg.base_url
        servers_cfg = Configuration.load_servers()

        # 连接多台 MCP 服务器
        mcp_client = MultiServerMCPClient(servers_cfg)
        tools = await mcp_client.get_tools()  # 工具对象列表
        logging.info(f"✅ 已加载 {len(tools)} 个 MCP 工具： {[t.name for t in tools]}")

        # 初始化大模型
        model = ChatQwQ(model="deepseek-v3", temperature=0)

        # 构造 LangGraph Agent
        agent = create_react_agent(model=model,
                                   tools=tools,
                                   prompt=prompt,
                                   checkpointer=checkpointer)

        ai_output = await agent.ainvoke(
            {"messages": [{"role": "user", "content": user_input}]},
            config
        )
        logging.info("agent_time_series正在发送数据")


        print("=== 代理工具调用过程追踪 ===")
        for i, message in enumerate(ai_output["messages"]):
            # 1. 打印AI消息，特别是包含 tool_calls 的消息
            if hasattr(message, 'tool_calls') and message.tool_calls:
                print(f"\n[{i}] 🤖 AI 决定调用工具:")
                for tool_call in message.tool_calls:
                    print(f"   工具名称: {tool_call['name']}")
                    print(f"   调用参数: {tool_call['args']}")
                    print(f"   调用ID: {tool_call['id']}")

            # 2. 打印工具返回的消息 (ToolMessage)
            # 根据你的输出，ToolMessage有 'name' 和 'content' 属性
            if hasattr(message, 'name') and message.name:
                print(f"\n[{i}] ⚙️ 工具 '{message.name}' 返回结果:")
                # 美化输出JSON或SQL
                content = message.content
                if content.strip().startswith(('{', '[', '```json')):
                    try:
                        # 尝试解析和美化JSON
                        import json
                        # 处理可能存在的代码块标记
                        if content.startswith('```json'):
                            content = content.strip('`').replace('json', '', 1).strip()
                        parsed_json = json.loads(content)
                        pretty_content = json.dumps(parsed_json, indent=2, ensure_ascii=False)
                        print(f"   返回内容:\n{pretty_content}")
                    except json.JSONDecodeError:
                        # 如果不是JSON，可能是SQL或其他格式
                        print(f"   返回内容: {content}")
                else:
                    print(f"   返回内容: {content}")

            # 也可以打印最终的AI回复
            if i == len(ai_output["messages"]) - 1 and hasattr(message, 'content') and message.content:
                print(f"\n[{i}] ✅ 最终回复: {message.content}")

        with open(file_folder_path / "chart_option.json", "r", encoding = "utf-8") as f:
            data = f.read()

        # 将结果打包成固定格式
        result = combine_info_to_post(await get_summary(user_input, ai_output['messages'][-1].content, data, agent), json.loads(data))

        # 给主智能体发送post请求
        update_insight_result(agent_type="agent_time_series", updated_result=result)
        logging.info("agent_time_series发送数据结束")

        return ai_output['messages'][-1].content

    elif mode == "non-agent":
        print(f'💡子智能体 agent_map 通过MCP 收到用户请求：{user_input}')
        inform_insight_thinking("agent_time_series", f"用户想{user_input}")
        key_template_json = await query_processing(user_input)
        print(key_template_json)
        key_template = json.loads(key_template_json)
        if key_template["query_type"] <= 2:

            max_retries = 3
            data = ""
            for attempt in range(max_retries):
                try:
                    sql = await text_to_sql(user_input, key_template_json)
                    data = await sql_inquire(sql)
                    print(f'💡子智能体 agent_time_series 根据用户输入生成了数据查询语句 {sql}')
                    inform_insight_thinking("agent_time_series",f'生成了数据查询语句 {sql[:20]}...')
                    print(f'💡子智能体 agent_time_series 根据数据查询语句得到查询结果 {data}')
                    inform_insight_thinking("agent_time_series",f'查询到了结果 {data[:20]}...')
                    break
                except Exception as e:
                    print(f"text2sql：尝试 {attempt + 1}/{max_retries} 失败: {e}")
                    if attempt >= max_retries - 1:
                        print("已达到最大重试次数，text2sql失败")
                        raise  "Failure in Text2Sql"
            data_array = json.loads(data)
            d_k, v_k = "日期", "default"
            res = []
            if data_array and "日期" in data_array[0].keys():
                for key in data_array[0].keys():
                    if key != "日期":
                        v_k = key
                for node in data_array:
                    res.append({"time": node[d_k], "value": node[v_k]})
            if not res: res = data_array

            predict_data = []
            if key_template["query_type"] == 1:
                predict_data = await predict_future(res, key_template["pred_day_length"])

            max_retries = 3
            chart_option = ""
            for attempt in range(max_retries):
                try:
                    chart_option = await visualize_line_chart(res, predict_data, key_template_json)
                    break
                except Exception as e:
                    print(f"visualize_line_chart：尝试 {attempt + 1}/{max_retries} 失败: {e}")
                    if attempt >= max_retries - 1:
                        print("已达到最大重试次数，visualize_line_chart失败")
                        raise "Failure in visualize_line_chart"

            answer, summary = await summary_factory(user_input, res, predict_data, key_template_json)

            print(f'💡子智能体 agent_time_series 根据查询结果生成了数据总结报告 {summary[:20]}...')
            inform_insight_thinking("agent_time_series", f'生成了数据总结报告 {summary[:20]}...')

            # 将结果打包成固定格式
            result = combine_info_to_post(summary, json.loads(chart_option))
            logging.info("agent_time_series正在发送数据")
            # 给主智能体发送post请求
            update_insight_result(agent_type="agent_time_series", updated_result=result)
            logging.info("agent_time_series发送数据结束")

            return answer
        elif key_template["query_type"] == 3:
            max_retries = 3
            data = ""
            for attempt in range(max_retries):
                try:
                    sql = await sankey_to_sql(user_input)
                    data = await sql_inquire(sql)
                    break
                except Exception as e:
                    print(f"text2sql：尝试 {attempt + 1}/{max_retries} 失败: {e}")
                    if attempt >= max_retries - 1:
                        print("已达到最大重试次数，text2sql失败")
                        raise "Failure in Text2Sql"

            max_retries = 3
            chart_option = ""
            for attempt in range(max_retries):
                try:
                    chart_option = await visualize_sankey_chart(data)
                    break
                except Exception as e:
                    print(f"visualize_line_chart：尝试 {attempt + 1}/{max_retries} 失败: {e}")
                    if attempt >= max_retries - 1:
                        print("已达到最大重试次数，visualize_line_chart失败")
                        raise "Failure in visualize_line_chart"
            answer, summary = await summary_factory(user_input, data, [], key_template_json)
            # 将结果打包成固定格式
            result = combine_info_to_post(summary, json.loads(chart_option))
            logging.info("agent_time_series正在发送数据")
            # 给主智能体发送post请求
            update_insight_result(agent_type="agent_time_series", updated_result=result)
            logging.info("agent_time_series发送数据结束")

            return answer
        else: raise "Non-Type Question"


    else:
        raise "Unsupported Mode"


if __name__ == "__main__":
    # logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
    # asyncio.run(respond_post("请帮我分析北京市2025年7月3日到2025年7月16日每日总家宽条数的变化"))


    mcp.run(transport='streamable-http')

    # logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
    # asyncio.run(agent_time_series("帮我分析20250717到20250730间北京市每日总家宽条数的变化"))